//
//  RBTree.swift
//  数据结构与算法
//
//  Created by Rowling on 2024/3/5.
//

import Foundation
var RED = false
var BLACK = true
class RBNode<T> : Node<T> {
    var color : Bool
    init(element : T, color: Bool) {
        self.color = color
        super.init(element: element)
    }
}

class RBTree<T : Comparable> : BinarySearchTree<T> {
    @discardableResult
    func color(node : Node<T>?, color : Bool) -> Node<T>? {
        guard let node = node as? RBNode else {return node}
        node.color = color
        return node
    }
    
    @discardableResult
    func red(_ node : Node<T>?) -> Node<T>? {
        return color(node: node, color: RED)
    }
    
    @discardableResult
    func black(_ node : Node<T>?) -> Node<T>? {
        return color(node: node, color: BLACK)
    }
    
    @discardableResult
    func colorOf(_ node : Node<T>?) -> Bool {
        guard let node = node as? RBNode else {return BLACK}
        return node.color
    }

    @discardableResult
    func isBlack(_ node : Node<T>?) -> Bool {
        return colorOf(node) == BLACK
    }
    
    @discardableResult
    func isRed(_ node : Node<T>?) -> Bool {
        return colorOf(node) == RED
    }
    
    override func createNode(element: T, left: Node<T>? = nil, right: Node<T>? = nil, parent: Node<T>? = nil) -> Node<T> {
        return RBNode(element: element, color: RED)
    }
    
    func leftRotation(node : Node<T>) {
        guard let right = node.right else {return}
        if node === root {
            root = right
        } else if let p = node.parent {
            if node.isLeft {
                p.left = right
            } else if node.isRight {
                p.right = right
            }
        }
        right.parent = node.parent
        node.parent = right
        node.right = right.left
        right.left?.parent = node
        right.left = node
    }
    
    func rightRotation(node : Node<T>) {
        guard let left = node.left else {return}
        if node === root {
            root = left
        } else if let p = node.parent {
            if p.left === node {
                p.left = left
            } else if p.right === node {
                p.right = left
            }
        }
        left.parent = node.parent
        node.parent = left
        node.left = left.right
        left.right?.parent = node
        left.right = node
    }
    
    override func afterAdd(node: Node<T>) {
        guard let parent = node.parent else {
            // 添加节点是跟节点
            _ = black(node)
            return
        }
        
        if isBlack(parent) {
            // 父节点是黑色，不用做处理
            return
        }
        
        let uncle = parent.sibling()
        let grand = parent.parent
        if isRed(uncle) {
            // 叔父是红色 --- 上溢
            black(parent)
            black(uncle)
            if let grand = red(grand) {
                afterAdd(node: grand)
            }
            return
        }
        guard let grand = grand else {return}
        red(grand)
        // 叔父节点不是红色
        if parent.isLeft {
            if node.isLeft {
                // LL
                rightRotation(node: grand)
                black(parent)
            } else {
                // LR
                leftRotation(node: parent)
                rightRotation(node: grand)
                black(node)
            }
        } else if parent.isRight {
            if node.isRight {
                // RR
                leftRotation(node: grand)
                black(parent)
            } else {
                // RL
                rightRotation(node: parent)
                leftRotation(node: grand)
                black(node)
            }
        }
    }
    
    override func afterRemove(node: Node<T>, replacement : Node<T>?) {
        // 如果删除的是红色直接干掉
        if isRed(node) {return}
        // 用于取代node子节点是红色
        if isRed(replacement) {
            black(replacement)
            return
        }
    
        
        guard let parent = node.parent else {
            // 删除的是根节点
            return
        }
        // 获取node的兄弟节点
        var isLeft = parent.left == nil
        if parent.left === node {
            isLeft = true
        } else if parent.right === node {
            isLeft = false
        }
        var sibling = isLeft ? parent.right : parent.left
//         node.sibling()
//        guard var sibling = sibling else {return}
        if isLeft { // 被删除的结点在左边
            if isRed(sibling) { // 兄弟节点是红色
                black(sibling)
                red(parent)
                leftRotation(node: parent)
                // 重新获取一下
                sibling = parent.right
            }
            // 能到这里，兄弟节点必然是黑色
            if isBlack(sibling?.left), isBlack(sibling?.right) { // 兄弟节点没有一个红色子节点
                if isBlack(parent) {
                    red(sibling)
                    afterRemove(node: parent, replacement: nil)
                    return
                } else {
                    black(parent)
                    red(sibling)
                }
            } else {
                // 兄弟至少有一个红色子节点
                if isBlack(sibling?.right) {
                    rightRotation(node: sibling!)
                    sibling = parent.right
                }
                color(node: sibling, color: colorOf(parent))
                black(sibling?.right)
                black(parent)
                leftRotation(node: parent)
            }

        } else { // 被删除的结点在右边
            if isRed(sibling) { // 兄弟节点是红色
                black(sibling)
                red(parent)
                rightRotation(node: parent)
                // 重新获取一下
                sibling = parent.left
            }
            // 能到这里，兄弟节点必然是黑色
            if isBlack(sibling?.left), isBlack(sibling?.right) { // 兄弟节点没有一个红色子节点
                if isBlack(parent) {
                    red(sibling)
                    afterRemove(node: parent, replacement: nil)
                    return
                } else {
                    black(parent)
                    red(sibling)
                }
            } else {
                // 兄弟至少有一个红色子节点
                if isBlack(sibling?.left) {
//                    black(sibling?.right)
//                    red(sibling)
                    leftRotation(node: sibling!)
                    sibling = parent.left
                }
                color(node: sibling, color: colorOf(parent))
                black(sibling?.left)
                black(parent)
                rightRotation(node: parent)
            }
        }
    }
}
